//////////////////////////////////////////////////////////////////////////////////
//                                                                              //
//  This file is part of the buola project (https://code.google.com/p/buola/).  //
//                                                                              //
//  Copyright(c) 2007-2012 Xavi Gratal                                          //
//  gratal AT gmail DOT com                                                     //
//                                                                              //
//  Buola is free software: you can redistribute it and/or modify               //
//  it under the terms of the GNU General Public License as published by        //
//  the Free Software Foundation, either version 3 of the License, or           //
//  (at your option) any later version.                                         //
//                                                                              //
//  Buola is distributed in the hope that it will be useful,                    //
//  but WITHOUT ANY WARRANTY; without even the implied warranty of              //
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               //
//  GNU General Public License for more details.                                //
//                                                                              //
//  You should have received a copy of the GNU General Public License           //
//  along with buola.  If not, see <http://www.gnu.org/licenses/>.              //
//                                                                              //
//////////////////////////////////////////////////////////////////////////////////

#include <buola/image/algorithm.h>
#include <buola/image/cimageplugin.h>
#include <buola/image/canyimage.h>

#include <png.h>

#pragma GCC diagnostic ignored "-Wclobbered"

namespace buola { namespace img {

class CPNGCodec : public CCodec
{
public:
    CPNGCodec();
    virtual ~CPNGCodec();

    void Read(io::PBufferStream pSrc,CAnyImage &pDst,EImageFormat pFmt,const CSize_i &pSize) final;
    void Write(const CAnyView &pSrc,io::PBufferStream pDst) final;

protected:
    CSize_i mSize;
    int mBitDepth;
    int mColorType;
    int mInterlaceType;

    png_structp mPNGPtr;
    png_infop mPNGInfo;
    png_infop mEndInfo;
};

class CPNGPlugin : public CImagePlugin
{
public:
    CPNGPlugin()
    {
        RegisterPlugin("png",this);
        RegisterPlugin("PNG",this);
    }

    CCodec *GetInstance()
    {
        return new CPNGCodec;
    }
};

static CPNGPlugin sPNGPlugin;

CPNGCodec::CPNGCodec()
{
}

CPNGCodec::~CPNGCodec()
{
}

void buola_png_read(png_structp pPNG,png_bytep pData,png_size_t pLength)
{
    io::CBufferStream *lStream=(io::CBufferStream*)png_get_io_ptr(pPNG);

    lStream->Read(pData,pLength);
}

void buola_png_write(png_structp pPNG,png_bytep pData,png_size_t pLength)
{
    io::CBufferStream *lStream=(io::CBufferStream*)png_get_io_ptr(pPNG);

    lStream->Write(pData,pLength);
}

void buola_png_flush(png_structp pPNG)
{
    io::CBufferStream *lStream=(io::CBufferStream*)png_get_io_ptr(pPNG);

    lStream->Flush();
}

void CPNGCodec::Read(io::PBufferStream pSrc,CAnyImage &pDst,EImageFormat pFmt,const CSize_i& /*pSize*/)
{
    png_structp lPNGPtr=png_create_read_struct(PNG_LIBPNG_VER_STRING,
                                                nullptr,nullptr,nullptr);

    if(!lPNGPtr)
    {
        throw XInternal("Couldn't open PNG file");
    }

    png_infop lPNGInfo=png_create_info_struct(lPNGPtr);
    png_infop lEndInfo=png_create_info_struct(lPNGPtr);

    if(!lPNGInfo||!lEndInfo)
    {
        png_destroy_read_struct(&lPNGPtr,nullptr,nullptr);
        throw XInternal("Couldn't open PNG file");
    }

    if(setjmp(png_jmpbuf(lPNGPtr)))
    {
        png_destroy_read_struct(&lPNGPtr,&lPNGInfo,&lEndInfo);
        throw XInternal("Error reading PNG file");
    }

    png_set_read_fn(lPNGPtr,pSrc.get(),buola_png_read);

    //bool lAlpha=false;

    png_read_info(lPNGPtr,lPNGInfo);
    png_uint_32 lPX,lPY;

    png_get_IHDR(lPNGPtr,lPNGInfo,&lPX,&lPY,
                    &mBitDepth,&mColorType,&mInterlaceType,nullptr,nullptr);

    mSize.x=lPX;
    mSize.y=lPY;

    if(mBitDepth==16)
        png_set_strip_16(lPNGPtr);

    if (mColorType==PNG_COLOR_TYPE_PALETTE)
    {
        png_set_palette_to_rgb(lPNGPtr);
        png_set_gray_to_rgb(lPNGPtr);
        png_set_filler(lPNGPtr,0xff,PNG_FILLER_AFTER);
        png_set_expand_gray_1_2_4_to_8(lPNGPtr);

        if(png_get_valid(lPNGPtr,lPNGInfo,PNG_INFO_tRNS))
        {
            png_set_tRNS_to_alpha(lPNGPtr);
            //lAlpha=true;
        }

        png_set_expand(lPNGPtr);
    }
    else if(mColorType==PNG_COLOR_TYPE_RGB)
    {
        png_set_filler(lPNGPtr,0xff,PNG_FILLER_AFTER);
    }
    else if(mColorType==PNG_COLOR_TYPE_RGB_ALPHA)
    {
        //lAlpha=true;
    }
    else if(mColorType==PNG_COLOR_TYPE_GRAY)
    {
        png_set_gray_to_rgb(lPNGPtr);
        png_set_filler(lPNGPtr,0,PNG_FILLER_AFTER);
        if(mBitDepth<8)
            png_set_expand_gray_1_2_4_to_8(lPNGPtr);
    }
    else if(mColorType==PNG_COLOR_TYPE_GRAY_ALPHA)
    {
        png_set_gray_to_rgb(lPNGPtr);
        //lAlpha=true;
    }

    png_read_update_info(lPNGPtr,lPNGInfo);

    int lBytesPerRow=png_get_rowbytes(lPNGPtr,lPNGInfo);

    unsigned char *mRowPointers[mSize.y];

    for(int i=0;i<mSize.y;i++)
    {
        mRowPointers[i]=alloc_array<unsigned char>(lBytesPerRow);
    }

    png_read_image(lPNGPtr,mRowPointers);
    png_read_end(lPNGPtr,lEndInfo);

    png_destroy_read_struct(&lPNGPtr,&lPNGInfo,&lEndInfo);

    unsigned char *lData=alloc_array<unsigned char>(lBytesPerRow*mSize.y+10);

    for(int i=0;i<mSize.y;i++)
    {
        memcpy(lData+i*lBytesPerRow,mRowPointers[i],lBytesPerRow);

        free(mRowPointers[i]);
    }

    CImage_rgba8 lImage(mSize);
    copy_pixels(img::view_rgba8(lData,mSize),lImage);

    pDst=std::move(lImage);
}

void CPNGCodec::Write(const img::CAnyView &pSrc,io::PBufferStream pDst)
{
    bool lSwap=false;
    
    int lDepth;
    int lPNGColorType;
    switch(pSrc.Format())
    {
    case img::FORMAT_BGR8:
        lSwap=true;
    case img::FORMAT_RGB8:
        lDepth=8;
        lPNGColorType=PNG_COLOR_TYPE_RGB;
        break;
    case img::FORMAT_BGRA8:
        lSwap=true;
    case img::FORMAT_RGBA8:
        lDepth=8;
        lPNGColorType=PNG_COLOR_TYPE_RGB_ALPHA;
        break;
    case img::FORMAT_GRAY8:
        lDepth=8;
        lPNGColorType=PNG_COLOR_TYPE_GRAY;
        break;
    case img::FORMAT_RGB16:
        lDepth=16;
        lPNGColorType=PNG_COLOR_TYPE_RGB;
        break;
    case img::FORMAT_RGBA16:
        lDepth=16;
        lPNGColorType=PNG_COLOR_TYPE_RGB_ALPHA;
        break;
    case img::FORMAT_GRAY16:
        lDepth=16;
        lPNGColorType=PNG_COLOR_TYPE_GRAY;
        break;
    default:
        throw XInvalid("invalid format writing PNG image");
    }
    
    png_structp lPNGPtr=png_create_write_struct(PNG_LIBPNG_VER_STRING,nullptr,nullptr,nullptr);
    if(!lPNGPtr) throw XInternal("Couldn't write PNG file");

    png_infop lPNGInfo=png_create_info_struct(lPNGPtr);

    if(!lPNGInfo)
    {
        png_destroy_write_struct(&lPNGPtr,nullptr);
        throw XInternal("Couldn't write PNG file");
    }

    if(setjmp(png_jmpbuf(lPNGPtr)))
    {
        png_destroy_write_struct(&lPNGPtr,&lPNGInfo);
        throw XInternal("Error writing PNG file");
    }

    png_set_write_fn(lPNGPtr,pDst.get(),buola_png_write,buola_png_flush);

    png_set_IHDR(lPNGPtr,lPNGInfo,pSrc.Width(),pSrc.Height(),lDepth,lPNGColorType,PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_DEFAULT,PNG_FILTER_TYPE_DEFAULT);

    png_write_info(lPNGPtr,lPNGInfo);

    if(lSwap)
        png_set_bgr(lPNGPtr);

    uint8_t *lRows[pSrc.Height()];
    for(int i=0;i<pSrc.Height();i++)
        lRows[i]=const_cast<uint8_t*>(pSrc.RawBytes(i));
    png_write_rows(lPNGPtr,lRows,pSrc.Height());
    png_write_end(lPNGPtr,lPNGInfo);

    png_destroy_write_struct(&lPNGPtr,&lPNGInfo);
}

/*namespace img*/ } /*namespace buola*/ }
